{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j2cz906V7d0X"
   },
   "source": [
    "# Training Apptronik Apollo using MuJoCo Warp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XpAzX-oCJcSJ"
   },
   "outputs": [],
   "source": [
    "# you may need some extra deps for this colab:\n",
    "# pip install \"jax[cuda12_local]\"\n",
    "# pip install playground\n",
    "# pip install matplotlib\n",
    "# if you want to run this on your local machine you can do like so:\n",
    "# pip install jupyter\n",
    "# jupyter notebook --NotebookApp.allow_origin='https://colab.research.google.com' --port=8888 --NotebookApp.port_retries=0 --no-browser\n",
    "\n",
    "import dataclasses\n",
    "import datetime\n",
    "import functools\n",
    "import os\n",
    "import time\n",
    "from typing import Any, Dict, Optional, Union\n",
    "\n",
    "import jax\n",
    "import mediapy as media\n",
    "import mujoco\n",
    "import numpy as np\n",
    "import warp as wp\n",
    "from etils import epath\n",
    "from jax import numpy as jp\n",
    "from ml_collections import config_dict\n",
    "from mujoco import mjx\n",
    "from mujoco_playground._src import mjx_env\n",
    "from mujoco_playground._src import reward\n",
    "from mujoco_playground._src.dm_control_suite import common\n",
    "from warp.jax_experimental.ffi import jax_callable\n",
    "\n",
    "import mujoco_warp as mjwarp\n",
    "\n",
    "# this ensures JAX embeds Warp kernels into its own computation graph:\n",
    "os.environ[\"XLA_FLAGS\"] = \"--xla_gpu_graph_min_graph_size=1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-BVYPgjk7xcm"
   },
   "outputs": [],
   "source": [
    "# We'll grab the Apptronik model from MuJoCo Menagerie, then remove some\n",
    "# MJX-specific changes that exist in the XML that MJWarp doesn't need\n",
    "# (such as explicit contacts, really tight ls_iterations etc.)\n",
    "\n",
    "mjx_env.ensure_menagerie_exists()\n",
    "\n",
    "contrib_xml_dir = epath.resource_path(\"mujoco_warp\").parent / \"contrib/xml\"\n",
    "apptronik_dir = mjx_env.EXTERNAL_DEPS_PATH / \"mujoco_menagerie/apptronik_apollo/\"\n",
    "\n",
    "! cp {contrib_xml_dir / 'apptronik_apollo.xml'} {apptronik_dir}\n",
    "! cp {contrib_xml_dir / 'scene.xml'} {apptronik_dir}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ixAOK05EkQNg"
   },
   "outputs": [],
   "source": [
    "# MJWarp is not yet fully conncected to JAX.  For now we will use the function\n",
    "# `mjwarp_step` below, which we wrap in Warp's handy `jax_callable` function\n",
    "# that converts Warp kernels to JAX operations.\n",
    "#\n",
    "# After we build a proper JAX wrapper for MJWarp, this code will disappear.\n",
    "\n",
    "NWORLD = 8192\n",
    "NCONMAX = 81920\n",
    "NJMAX = NCONMAX * 4\n",
    "\n",
    "xml_path = apptronik_dir / \"scene.xml\"\n",
    "mjm = mujoco.MjModel.from_xml_path(xml_path.as_posix())\n",
    "mjm.opt.iterations = 5\n",
    "mjm.opt.ls_iterations = 10\n",
    "mjd = mujoco.MjData(mjm)\n",
    "mujoco.mj_resetDataKeyframe(mjm, mjd, 0)\n",
    "mujoco.mj_forward(mjm, mjd)\n",
    "m = mjwarp.put_model(mjm)\n",
    "d = mjwarp.put_data(mjm, mjd, nworld=NWORLD, nconmax=NCONMAX, njmax=NJMAX)\n",
    "\n",
    "\n",
    "def mjwarp_step(\n",
    "  ctrl: wp.array(dtype=wp.float32, ndim=2),\n",
    "  qpos_in: wp.array(dtype=wp.float32, ndim=2),\n",
    "  qvel_in: wp.array(dtype=wp.float32, ndim=2),\n",
    "  qacc_warmstart_in: wp.array(dtype=wp.float32, ndim=2),\n",
    "  qpos_out: wp.array(dtype=wp.float32, ndim=2),\n",
    "  qvel_out: wp.array(dtype=wp.float32, ndim=2),\n",
    "  xpos_out: wp.array(dtype=wp.vec3, ndim=2),\n",
    "  xmat_out: wp.array(dtype=wp.mat33, ndim=2),\n",
    "  qacc_warmstart_out: wp.array(dtype=wp.float32, ndim=2),\n",
    "  subtree_com_out: wp.array(dtype=wp.vec3, ndim=2),\n",
    "  cvel_out: wp.array(dtype=wp.spatial_vector, ndim=2),\n",
    "  site_xpos_out: wp.array(dtype=wp.vec3, ndim=2),\n",
    "):\n",
    "  wp.copy(d.ctrl, ctrl)\n",
    "  wp.copy(d.qpos, qpos_in)\n",
    "  wp.copy(d.qvel, qvel_in)\n",
    "  wp.copy(d.qacc_warmstart, qacc_warmstart_in)\n",
    "\n",
    "  # TODO(team): remove this hard coding substeps\n",
    "  # ctrl_dt / sim_dt == 4\n",
    "  for i in range(4):\n",
    "    mjwarp.step(m, d)\n",
    "  wp.copy(qpos_out, d.qpos)\n",
    "  wp.copy(qvel_out, d.qvel)\n",
    "  wp.copy(xpos_out, d.xpos)\n",
    "  wp.copy(xmat_out, d.xmat)\n",
    "  wp.copy(qacc_warmstart_out, d.qacc_warmstart)\n",
    "  wp.copy(subtree_com_out, d.subtree_com)\n",
    "  wp.copy(cvel_out, d.cvel)\n",
    "  wp.copy(site_xpos_out, d.site_xpos)\n",
    "\n",
    "\n",
    "jax_mjwarp_step = jax_callable(\n",
    "  mjwarp_step,\n",
    "  num_outputs=8,\n",
    "  output_dims={\n",
    "    \"qpos_out\": (NWORLD, mjm.nq),\n",
    "    \"qvel_out\": (NWORLD, mjm.nv),\n",
    "    \"xpos_out\": (NWORLD, mjm.nbody, 3),\n",
    "    \"xmat_out\": (NWORLD, mjm.nbody, 3, 3),\n",
    "    \"qacc_warmstart_out\": (NWORLD, mjm.nv),\n",
    "    \"subtree_com_out\": (NWORLD, mjm.nbody, 3),\n",
    "    \"cvel_out\": (NWORLD, mjm.nbody, 6),\n",
    "    \"site_xpos_out\": (NWORLD, mjm.nsite, 3),\n",
    "  },\n",
    ")\n",
    "\n",
    "# the functions below allow us to call MJWarp step inside jax vmap:\n",
    "\n",
    "\n",
    "@jax.custom_batching.custom_vmap\n",
    "def step(d: mjx.Data):\n",
    "  return d\n",
    "\n",
    "\n",
    "@step.def_vmap\n",
    "def step_vmap_rule(axis_size, in_batched, d: mjx.Data):\n",
    "  if in_batched[0].ctrl:\n",
    "    assert d.ctrl.shape[0] == axis_size\n",
    "  else:\n",
    "    d = d.replace(ctrl=jp.tile(d.ctrl, (axis_size, 1)))\n",
    "  params = {f.name: None for f in dataclasses.fields(mjx.Data)}\n",
    "  params[\"ctrl\"] = True\n",
    "  params[\"qpos\"] = True\n",
    "  params[\"qvel\"] = True\n",
    "  params[\"xpos\"] = True\n",
    "  params[\"xmat\"] = True\n",
    "  params[\"qacc_warmstart\"] = True\n",
    "  params[\"subtree_com\"] = True\n",
    "  params[\"cvel\"] = True\n",
    "  params[\"site_xpos\"] = True\n",
    "  out_batched = mjx.Data(**params)\n",
    "\n",
    "  qpos, qvel, xpos, xmat, qacc_warmstart, subtree_com, cvel, site_xpos = jax_mjwarp_step(\n",
    "    d.ctrl, d.qpos, d.qvel, d.qacc_warmstart\n",
    "  )\n",
    "  d = d.replace(\n",
    "    qpos=qpos,\n",
    "    qvel=qvel,\n",
    "    xpos=xpos,\n",
    "    xmat=xmat,\n",
    "    qacc_warmstart=qacc_warmstart,\n",
    "    subtree_com=subtree_com,\n",
    "    cvel=cvel,\n",
    "    site_xpos=site_xpos,\n",
    "  )\n",
    "  return d, out_batched\n",
    "\n",
    "\n",
    "def init(qpos, ctrl) -> mjx.Data:\n",
    "  init_params = {f.name: None for f in dataclasses.fields(mjx.Data)}\n",
    "  init_params[\"qpos\"] = qpos\n",
    "  init_params[\"ctrl\"] = ctrl\n",
    "  init_params[\"qvel\"] = jp.zeros(m.nv)\n",
    "  init_params[\"xpos\"] = jp.zeros((m.nbody, 3))\n",
    "  init_params[\"xmat\"] = jp.tile(jp.zeros((3, 3)), (m.nbody, 1, 1))\n",
    "  init_params[\"qacc_warmstart\"] = jp.array(mjd.qacc_warmstart)\n",
    "  init_params[\"subtree_com\"] = jp.zeros((m.nbody, 3))\n",
    "  init_params[\"cvel\"] = jp.zeros((m.nbody, 6))\n",
    "  init_params[\"site_xpos\"] = jp.zeros((m.nsite, 3))\n",
    "  return mjx.Data(**init_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AaYQk2lu7Bus"
   },
   "outputs": [],
   "source": [
    "# An environment for training an Apptronik Apollo to walk.\n",
    "# This is the same format as environments in MuJoCo Playground.\n",
    "\n",
    "\n",
    "def default_config() -> config_dict.ConfigDict:\n",
    "  return config_dict.create(\n",
    "    ctrl_dt=0.02,\n",
    "    sim_dt=0.005,\n",
    "    episode_length=1000,\n",
    "    action_repeat=1,\n",
    "    action_scale=0.5,\n",
    "    soft_joint_pos_limit_factor=0.95,\n",
    "    reward_config=config_dict.create(\n",
    "      scales=config_dict.create(\n",
    "        # Tracking related rewards.\n",
    "        tracking_lin_vel=1.0,\n",
    "        tracking_ang_vel=0.75,\n",
    "        # Base related rewards.\n",
    "        ang_vel_xy=-0.15,\n",
    "        orientation=-2.0,\n",
    "        # Energy related rewards.\n",
    "        action_rate=0.0,\n",
    "        # Feet related rewards.\n",
    "        feet_air_time=2.0,\n",
    "        feet_slip=-0.25,\n",
    "        feet_phase=1.0,\n",
    "        # Other rewards.\n",
    "        termination=-5.0,\n",
    "        # Pose related rewards.\n",
    "        joint_deviation_knee=-0.1,\n",
    "        joint_deviation_hip=-0.25,\n",
    "        dof_pos_limits=-1.0,\n",
    "        pose=-0.1,\n",
    "      ),\n",
    "      tracking_sigma=0.25,\n",
    "      max_foot_height=0.15,\n",
    "      base_height_target=0.5,\n",
    "    ),\n",
    "    lin_vel_x=[1.0, 1.0],\n",
    "    lin_vel_y=[0.0, 0.0],\n",
    "    ang_vel_yaw=[0.0, 0.0],\n",
    "  )\n",
    "\n",
    "\n",
    "class Joystick(mjx_env.MjxEnv):\n",
    "  \"\"\"Track a joystick command.\"\"\"\n",
    "\n",
    "  def __init__(\n",
    "    self,\n",
    "    config: config_dict.ConfigDict = default_config(),\n",
    "    config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,\n",
    "  ):\n",
    "    super().__init__(config, config_overrides)\n",
    "    self._post_init()\n",
    "\n",
    "  def _post_init(self) -> None:\n",
    "    self._init_q = jp.array(mjm.keyframe(\"stand\").qpos)\n",
    "    self._default_pose = self._init_q[7:]\n",
    "\n",
    "    # Note: First joint is freejoint.\n",
    "    self._lowers, self._uppers = self.mj_model.jnt_range[1:].T\n",
    "    c = (self._lowers + self._uppers) / 2\n",
    "    r = self._uppers - self._lowers\n",
    "    self._soft_lowers = c - 0.5 * r * self._config.soft_joint_pos_limit_factor\n",
    "    self._soft_uppers = c + 0.5 * r * self._config.soft_joint_pos_limit_factor\n",
    "\n",
    "    hip_joints = [\"l_hip_ie\", \"l_hip_aa\", \"r_hip_ie\", \"r_hip_aa\"]\n",
    "    hip_indices = [mjm.joint(j).qposadr - 7 for j in hip_joints]\n",
    "    self._hip_indices = jp.array(hip_indices)\n",
    "\n",
    "    knee_joints = [\"l_knee_fe\", \"r_knee_fe\"]\n",
    "    knee_indices = [mjm.joint(j).qposadr - 7 for j in knee_joints]\n",
    "    self._knee_indices = jp.array(knee_indices)\n",
    "\n",
    "    self._head_body_id = mjm.body(\"neck_pitch_link\").id\n",
    "    self._torso_id = mjm.body(\"torso_link\").id\n",
    "\n",
    "    feet_sites = [\"l_foot_fr\", \"l_foot_br\", \"l_foot_fl\", \"l_foot_bl\"]\n",
    "    feet_sites += [\"r_foot_fr\", \"r_foot_br\", \"r_foot_fl\", \"r_foot_bl\"]\n",
    "    feet_site_ids = [mjm.site(s).id for s in feet_sites]\n",
    "    self._feet_site_id = jp.array(feet_site_ids)\n",
    "    self._feet_contact_z = 0.003\n",
    "\n",
    "    self._floor_geom_id = mjm.geom(\"floor\").id\n",
    "\n",
    "  def reset(self, rng: jax.Array) -> mjx_env.State:\n",
    "    qpos = self._init_q\n",
    "\n",
    "    data = init(qpos=qpos, ctrl=qpos[7:])\n",
    "\n",
    "    # Phase, freq=U(1.0, 1.5)\n",
    "    rng, key = jax.random.split(rng)\n",
    "    gait_freq = jax.random.uniform(key, (1,), minval=1.25, maxval=1.5)\n",
    "    phase_dt = 2 * jp.pi * self.dt * gait_freq\n",
    "    phase = jp.array([0, jp.pi])\n",
    "\n",
    "    rng, cmd_rng = jax.random.split(rng)\n",
    "    cmd = self.sample_command(cmd_rng)\n",
    "\n",
    "    info = {\n",
    "      \"rng\": rng,\n",
    "      \"step\": 0,\n",
    "      \"command\": cmd,\n",
    "      \"last_act\": jp.zeros(mjm.nu),\n",
    "      \"last_last_act\": jp.zeros(mjm.nu),\n",
    "      \"motor_targets\": jp.zeros(mjm.nu),\n",
    "      \"feet_air_time\": jp.zeros(2),\n",
    "      \"last_contact\": jp.zeros(2, dtype=bool),\n",
    "      # Phase related.\n",
    "      \"phase_dt\": phase_dt,\n",
    "      \"phase\": phase,\n",
    "    }\n",
    "\n",
    "    metrics = {}\n",
    "    for k in self._config.reward_config.scales.keys():\n",
    "      metrics[f\"reward/{k}\"] = jp.zeros(())\n",
    "\n",
    "    contact = data.site_xpos[self._feet_site_id] < self._feet_contact_z\n",
    "    contact = jp.array([contact[0:4].any(), contact[4:8].any()])\n",
    "    obs = self._get_obs(data, info, contact)\n",
    "    reward, done = jp.zeros(2)\n",
    "    return mjx_env.State(data, obs, reward, done, metrics, info)\n",
    "\n",
    "  def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:\n",
    "    state.info[\"rng\"], _ = jax.random.split(state.info[\"rng\"], 2)\n",
    "\n",
    "    ctrl = self._default_pose + action * self._config.action_scale\n",
    "    data = state.data\n",
    "    data = data.replace(ctrl=ctrl)\n",
    "    data = step(data)\n",
    "    state.info[\"motor_targets\"] = ctrl\n",
    "\n",
    "    contact = data.site_xpos[self._feet_site_id, 2] < self._feet_contact_z\n",
    "    contact = jp.array([contact[0:4].any(), contact[4:8].any()])\n",
    "    contact_filt = contact | state.info[\"last_contact\"]\n",
    "    first_contact = (state.info[\"feet_air_time\"] > 0.0) * contact_filt\n",
    "    state.info[\"feet_air_time\"] += self.dt\n",
    "\n",
    "    obs = self._get_obs(data, state.info, contact)\n",
    "    done = self._get_termination(data)\n",
    "\n",
    "    rewards = self._get_reward(data, action, state.info, state.metrics, done, first_contact, contact)\n",
    "    rewards = {k: v * self._config.reward_config.scales[k] for k, v in rewards.items()}\n",
    "    reward = sum(rewards.values()) * self.dt\n",
    "\n",
    "    state.info[\"step\"] += 1\n",
    "    phase_tp1 = state.info[\"phase\"] + state.info[\"phase_dt\"]\n",
    "    state.info[\"phase\"] = jp.fmod(phase_tp1 + jp.pi, 2 * jp.pi) - jp.pi\n",
    "    state.info[\"last_last_act\"] = state.info[\"last_act\"]\n",
    "    state.info[\"last_act\"] = action\n",
    "    state.info[\"rng\"], cmd_rng = jax.random.split(state.info[\"rng\"])\n",
    "    state.info[\"command\"] = jp.where(\n",
    "      state.info[\"step\"] > 500,\n",
    "      self.sample_command(cmd_rng),\n",
    "      state.info[\"command\"],\n",
    "    )\n",
    "    state.info[\"step\"] = jp.where(\n",
    "      done | (state.info[\"step\"] > 500),\n",
    "      0,\n",
    "      state.info[\"step\"],\n",
    "    )\n",
    "    state.info[\"feet_air_time\"] *= ~contact\n",
    "    state.info[\"last_contact\"] = contact\n",
    "    for k, v in rewards.items():\n",
    "      state.metrics[f\"reward/{k}\"] = v\n",
    "\n",
    "    done = done.astype(reward.dtype)\n",
    "    state = state.replace(data=data, obs=obs, reward=reward, done=done)\n",
    "    return state\n",
    "\n",
    "  def _get_termination(self, data: mjx.Data) -> jax.Array:\n",
    "    fall_termination = data.xpos[self._head_body_id, 2] < 1.0\n",
    "    return fall_termination\n",
    "\n",
    "  def _get_obs(self, data: mjx.Data, info: dict[str, Any], contact: jax.Array) -> mjx_env.Observation:\n",
    "    cos = jp.cos(info[\"phase\"])\n",
    "    sin = jp.sin(info[\"phase\"])\n",
    "    phase = jp.concatenate([cos, sin])\n",
    "\n",
    "    return jp.hstack(\n",
    "      [\n",
    "        data.qpos,\n",
    "        data.qvel,\n",
    "        data.cvel.ravel(),\n",
    "        data.xpos.ravel(),\n",
    "        data.xmat.ravel(),\n",
    "        phase,\n",
    "        info[\"command\"],\n",
    "        info[\"last_act\"],\n",
    "        info[\"feet_air_time\"],\n",
    "      ]\n",
    "    )\n",
    "\n",
    "  def _get_reward(\n",
    "    self,\n",
    "    data: mjx.Data,\n",
    "    action: jax.Array,\n",
    "    info: dict[str, Any],\n",
    "    metrics: dict[str, Any],\n",
    "    done: jax.Array,\n",
    "    first_contact: jax.Array,\n",
    "    contact: jax.Array,\n",
    "  ) -> dict[str, jax.Array]:\n",
    "    del metrics  # Unused.\n",
    "    return {\n",
    "      # Tracking rewards.\n",
    "      \"tracking_lin_vel\": self._reward_tracking_lin_vel(info[\"command\"], self._get_global_linvel(data, self._torso_id)),\n",
    "      \"tracking_ang_vel\": self._reward_tracking_ang_vel(info[\"command\"], self._get_global_angvel(data, self._torso_id)),\n",
    "      # Base-related rewards.\n",
    "      \"ang_vel_xy\": self._cost_ang_vel_xy(self._get_global_angvel(data, self._torso_id)),\n",
    "      \"orientation\": self._cost_orientation(self._get_z_frame(data, self._torso_id)),\n",
    "      # Energy related rewards.\n",
    "      \"action_rate\": self._cost_action_rate(action, info[\"last_act\"], info[\"last_last_act\"]),\n",
    "      # Feet related rewards.\n",
    "      \"feet_slip\": self._cost_feet_slip(data, contact, info),\n",
    "      \"feet_air_time\": self._reward_feet_air_time(info[\"feet_air_time\"], first_contact, info[\"command\"]),\n",
    "      \"feet_phase\": self._reward_feet_phase(\n",
    "        data,\n",
    "        info[\"phase\"],\n",
    "        self._config.reward_config.max_foot_height,\n",
    "        info[\"command\"],\n",
    "      ),\n",
    "      # Pose related rewards.\n",
    "      \"joint_deviation_hip\": self._cost_joint_deviation_hip(data.qpos[7:], info[\"command\"]),\n",
    "      \"joint_deviation_knee\": self._cost_joint_deviation_knee(data.qpos[7:]),\n",
    "      \"dof_pos_limits\": self._cost_joint_pos_limits(data.qpos[7:]),\n",
    "      \"pose\": self._cost_pose(data.qpos[7:]),\n",
    "      # Other rewards.\n",
    "      \"termination\": self._cost_termination(done),\n",
    "    }\n",
    "\n",
    "  def _get_global_angvel(self, data: mjx.Data, bodyid: int):\n",
    "    return data.cvel[bodyid, :3]\n",
    "\n",
    "  def _get_global_linvel(self, data: mjx.Data, bodyid: int):\n",
    "    offset = data.xpos[bodyid] - data.subtree_com[mjm.body_rootid[bodyid]]\n",
    "    xang = data.cvel[bodyid, :3]\n",
    "    xvel = data.cvel[bodyid, 3:] + jp.cross(offset, xang)\n",
    "    return xvel\n",
    "\n",
    "  def _get_z_frame(self, data: mjx.Data, bodyid: int):\n",
    "    return data.xmat[bodyid, :, 2]\n",
    "\n",
    "  # Tracking rewards.\n",
    "\n",
    "  def _reward_tracking_lin_vel(\n",
    "    self,\n",
    "    commands: jax.Array,\n",
    "    local_vel: jax.Array,\n",
    "  ) -> jax.Array:\n",
    "    lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))\n",
    "    return jp.exp(-lin_vel_error / self._config.reward_config.tracking_sigma)\n",
    "\n",
    "  def _reward_tracking_ang_vel(\n",
    "    self,\n",
    "    commands: jax.Array,\n",
    "    ang_vel: jax.Array,\n",
    "  ) -> jax.Array:\n",
    "    ang_vel_error = jp.square(commands[2] - ang_vel[2])\n",
    "    return jp.exp(-ang_vel_error / self._config.reward_config.tracking_sigma)\n",
    "\n",
    "  # Base-related rewards.\n",
    "\n",
    "  def _cost_ang_vel_xy(self, global_angvel_torso: jax.Array) -> jax.Array:\n",
    "    return jp.sum(jp.square(global_angvel_torso[:2]))\n",
    "\n",
    "  def _cost_orientation(self, torso_zaxis: jax.Array) -> jax.Array:\n",
    "    return jp.sum(jp.square(torso_zaxis - jp.array([0.0, 0.0, 1.0])))\n",
    "\n",
    "  def _cost_base_height(self, base_height: jax.Array) -> jax.Array:\n",
    "    return jp.square(base_height - self._config.reward_config.base_height_target)\n",
    "\n",
    "  # Energy related rewards.\n",
    "\n",
    "  def _cost_action_rate(self, act: jax.Array, last_act: jax.Array, last_last_act: jax.Array) -> jax.Array:\n",
    "    del last_last_act  # Unused.\n",
    "    return jp.sum(jp.square(act - last_act))\n",
    "\n",
    "  # Feet related rewards.\n",
    "\n",
    "  def _cost_feet_slip(self, data: mjx.Data, contact: jax.Array, info: dict[str, Any]) -> jax.Array:\n",
    "    del info  # Unused.\n",
    "    body_vel = self._get_global_linvel(data, self._torso_id)[:2]\n",
    "    reward = jp.sum(jp.linalg.norm(body_vel, axis=-1) * contact)\n",
    "    return reward\n",
    "\n",
    "  def _reward_feet_air_time(\n",
    "    self,\n",
    "    air_time: jax.Array,\n",
    "    first_contact: jax.Array,\n",
    "    commands: jax.Array,\n",
    "    threshold_min: float = 0.2,\n",
    "    threshold_max: float = 0.5,\n",
    "  ) -> jax.Array:\n",
    "    del commands  # Unused.\n",
    "    air_time = (air_time - threshold_min) * first_contact\n",
    "    air_time = jp.clip(air_time, max=threshold_max - threshold_min)\n",
    "    reward = jp.sum(air_time)\n",
    "    return reward\n",
    "\n",
    "  def get_rz(phi: Union[jax.Array, float], swing_height: Union[jax.Array, float] = 0.08) -> jax.Array:\n",
    "    def cubic_bezier_interpolation(y_start, y_end, x):\n",
    "      y_diff = y_end - y_start\n",
    "      bezier = x**3 + 3 * (x**2 * (1 - x))\n",
    "      return y_start + y_diff * bezier\n",
    "\n",
    "    x = (phi + jp.pi) / (2 * jp.pi)\n",
    "    stance = cubic_bezier_interpolation(0, swing_height, 2 * x)\n",
    "    swing = cubic_bezier_interpolation(swing_height, 0, 2 * x - 1)\n",
    "    return jp.where(x <= 0.5, stance, swing)\n",
    "\n",
    "  def _reward_feet_phase(\n",
    "    self,\n",
    "    data: mjx.Data,\n",
    "    phase: jax.Array,\n",
    "    foot_height: jax.Array,\n",
    "    command: jax.Array,\n",
    "  ) -> jax.Array:\n",
    "    # Reward for tracking the desired foot height.\n",
    "    foot_pos = data.site_xpos[self._feet_site_id]\n",
    "    foot_pos = jp.array([jp.mean(foot_pos[0:4], axis=0), jp.mean(foot_pos[4:8], axis=0)])\n",
    "    foot_z = foot_pos[..., -1]\n",
    "    rz = Joystick.get_rz(phase, swing_height=foot_height)\n",
    "    error = jp.sum(jp.square(foot_z - rz))\n",
    "    reward = jp.exp(-error / 0.01)\n",
    "    body_linvel = self._get_global_linvel(data, self._torso_id)[:2]\n",
    "    body_angvel = self._get_global_angvel(data, self._torso_id)[2]\n",
    "    linvel_mask = jp.logical_or(\n",
    "      jp.linalg.norm(body_linvel) > 0.1,\n",
    "      jp.abs(body_angvel) > 0.1,\n",
    "    )\n",
    "    mask = jp.logical_or(linvel_mask, jp.linalg.norm(command) > 0.01)\n",
    "    reward *= mask\n",
    "    return reward\n",
    "\n",
    "  # Pose-related rewards.\n",
    "\n",
    "  def _cost_joint_deviation_hip(self, qpos: jax.Array, cmd: jax.Array) -> jax.Array:\n",
    "    error = qpos[self._hip_indices] - self._default_pose[self._hip_indices]\n",
    "    # Allow roll deviation when lateral velocity is high.\n",
    "    weight = jp.where(\n",
    "      cmd[1] > 0.1,\n",
    "      jp.array([0.0, 1.0, 0.0, 1.0]),\n",
    "      jp.array([1.0, 1.0, 1.0, 1.0]),\n",
    "    )\n",
    "    cost = jp.sum(jp.abs(error) * weight)\n",
    "    return cost\n",
    "\n",
    "  def _cost_joint_deviation_knee(self, qpos: jax.Array) -> jax.Array:\n",
    "    error = qpos[self._knee_indices] - self._default_pose[self._knee_indices]\n",
    "    return jp.sum(jp.abs(error))\n",
    "\n",
    "  def _cost_joint_pos_limits(self, qpos: jax.Array) -> jax.Array:\n",
    "    out_of_limits = -jp.clip(qpos - self._soft_lowers, None, 0.0)\n",
    "    out_of_limits += jp.clip(qpos - self._soft_uppers, 0.0, None)\n",
    "    return jp.sum(out_of_limits)\n",
    "\n",
    "  def _cost_pose(self, qpos: jax.Array) -> jax.Array:\n",
    "    return jp.sum(jp.square(qpos - self._default_pose))\n",
    "\n",
    "  # Other rewards.\n",
    "\n",
    "  def _cost_termination(self, done: jax.Array) -> jax.Array:\n",
    "    return done\n",
    "\n",
    "  def sample_command(self, rng: jax.Array) -> jax.Array:\n",
    "    rng1, rng2, rng3, rng4 = jax.random.split(rng, 4)\n",
    "\n",
    "    lin_vel_x = jax.random.uniform(rng1, minval=self._config.lin_vel_x[0], maxval=self._config.lin_vel_x[1])\n",
    "    lin_vel_y = jax.random.uniform(rng2, minval=self._config.lin_vel_y[0], maxval=self._config.lin_vel_y[1])\n",
    "    ang_vel_yaw = jax.random.uniform(\n",
    "      rng3,\n",
    "      minval=self._config.ang_vel_yaw[0],\n",
    "      maxval=self._config.ang_vel_yaw[1],\n",
    "    )\n",
    "\n",
    "    return jp.hstack([lin_vel_x, lin_vel_y, ang_vel_yaw])\n",
    "\n",
    "  @property\n",
    "  def xml_path(self) -> str:\n",
    "    return xml_path.as_posix()\n",
    "\n",
    "  @property\n",
    "  def action_size(self) -> int:\n",
    "    return mjm.nu\n",
    "\n",
    "  @property\n",
    "  def mj_model(self) -> mujoco.MjModel:\n",
    "    return mjm\n",
    "\n",
    "  @property\n",
    "  def mjx_model(self) -> mjx.Model:\n",
    "    return None  # unused"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aoWmBHXiU-HW"
   },
   "outputs": [],
   "source": [
    "# Train the environment using Brax PPO\n",
    "\n",
    "from brax.training.agents.ppo import networks as ppo_networks\n",
    "from brax.training.agents.ppo import train as ppo\n",
    "from IPython.display import HTML\n",
    "from IPython.display import clear_output\n",
    "from matplotlib import pyplot as plt\n",
    "from mujoco_playground import wrapper\n",
    "from mujoco_playground.config import locomotion_params\n",
    "\n",
    "ppo_params = config_dict.create(\n",
    "  num_timesteps=50_000_000,\n",
    "  num_evals=5,\n",
    "  reward_scaling=1.0,\n",
    "  clipping_epsilon=0.2,\n",
    "  episode_length=1000,\n",
    "  normalize_observations=True,\n",
    "  action_repeat=1,\n",
    "  unroll_length=20,\n",
    "  num_minibatches=32,\n",
    "  num_updates_per_batch=4,\n",
    "  discounting=0.97,\n",
    "  learning_rate=3e-4,\n",
    "  entropy_cost=0.005,\n",
    "  num_envs=NWORLD,\n",
    "  num_eval_envs=NWORLD,\n",
    "  batch_size=256,\n",
    "  max_grad_norm=1.0,\n",
    "  # network_factory=config_dict.create(\n",
    "  #   policy_hidden_layer_sizes=(512, 256, 128),\n",
    "  #   value_hidden_layer_sizes=(512, 256, 128),\n",
    "  # ),\n",
    ")\n",
    "\n",
    "x_data, y_data, y_dataerr = [], [], []\n",
    "times = [datetime.datetime.now()]\n",
    "\n",
    "\n",
    "def progress(num_steps, metrics):\n",
    "  times.append(datetime.datetime.now())\n",
    "  x_data.append(num_steps)\n",
    "  y_data.append(metrics[\"eval/episode_reward\"])\n",
    "  y_dataerr.append(metrics[\"eval/episode_reward_std\"])\n",
    "\n",
    "  plt.xlim([0, ppo_params[\"num_timesteps\"] * 1.25])\n",
    "  plt.ylim([0, 30])\n",
    "  plt.xlabel(\"# environment steps\")\n",
    "  plt.ylabel(\"reward per episode\")\n",
    "  plt.title(f\"y={y_data[-1]:.3f}\")\n",
    "  plt.errorbar(x_data, y_data, yerr=y_dataerr, color=\"blue\")\n",
    "\n",
    "  display(plt.gcf())\n",
    "  clear_output(wait=True)\n",
    "\n",
    "\n",
    "ppo_training_params = dict(ppo_params)\n",
    "network_factory = ppo_networks.make_ppo_networks\n",
    "if \"network_factory\" in ppo_params:\n",
    "  del ppo_training_params[\"network_factory\"]\n",
    "  network_factory = functools.partial(ppo_networks.make_ppo_networks, **ppo_params.network_factory)\n",
    "\n",
    "train_fn = functools.partial(\n",
    "  ppo.train,\n",
    "  **dict(ppo_training_params),\n",
    "  network_factory=network_factory,\n",
    "  progress_fn=progress,\n",
    ")\n",
    "\n",
    "env = Joystick()\n",
    "\n",
    "make_inference_fn, params, metrics = train_fn(\n",
    "  environment=env,\n",
    "  wrap_env_fn=wrapper.wrap_for_brax_training,\n",
    ")\n",
    "\n",
    "print(f\"time to jit: {times[1] - times[0]}\")\n",
    "print(f\"time to train: {times[-1] - times[1]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gGRhELoRjtQs"
   },
   "outputs": [],
   "source": [
    "rng = jax.random.PRNGKey(0)\n",
    "\n",
    "jit_reset = jax.jit(env.reset)\n",
    "\n",
    "\n",
    "def unroll(state):\n",
    "  inference_fn = make_inference_fn(params, deterministic=True)\n",
    "  rng = jax.random.PRNGKey(0)\n",
    "\n",
    "  def single_step(state, _):\n",
    "    action, _ = inference_fn(state.obs, rng)\n",
    "    action = jp.tile(action, (NWORLD, 1))\n",
    "    state = jax.tree.map(lambda x: jp.tile(x, (NWORLD,) + (1,) * len(x.shape)), state)\n",
    "    state = jax.vmap(env.step)(state, action)\n",
    "    state = jax.tree.map(lambda x: x[0], state)\n",
    "\n",
    "    return state, state\n",
    "\n",
    "  _, states = jax.lax.scan(single_step, state, length=1000)\n",
    "\n",
    "  return states\n",
    "\n",
    "\n",
    "rollout = jax.jit(unroll)(jit_reset(rng))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_kj9fCiDn6j6"
   },
   "outputs": [],
   "source": [
    "rollout_arr = [jax.tree.map(lambda x, i=i: x[i], rollout) for i in range(400)]\n",
    "frames = env.render(rollout_arr, camera=\"track\", width=640, height=480)\n",
    "media.show_video(frames, fps=1.0 / env.dt)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
